#include <stdlib.h>
#include <string.h>

#include "tdp2.hpp"
#include "tdp2_rc5_config.hpp"
#include "tdp2_rc5_blocksort.hpp"
#include "tdp2_rc5_sm.hpp"
#include "tdp2_rc5.hpp"
#include "tdp2_rc.hpp"

using namespace TDP;
using namespace TDP::Compression;
using namespace TDP::Compression::RangeCoder;
using namespace TDP::Compression::BlockSort;
using namespace TDP::Compression::RC5;
using namespace TDP::Compression::RC5::StructuredModel;

static const void *RC5BlockSort(const void *input, void *tempspace, UInt32 count, UInt32 *pindex, UInt32 *ftab)
{
	UInt32 i, pi, n;
	UInt32 *arr1;
	UInt32 *arr2;
	const UInt8 *inc;
	UInt8 *outc;

	// Prep blocksort
	arr1 = reinterpret_cast<UInt32*>(tempspace);
	arr2 = arr1 + count;

	memcpy(arr2, input, count);

	// Sort
	TDP::Compression::BlockSort::Sort(ftab, count, arr1, arr2);

	// Dump BWT output.  arr2 can be progressively overwritten
	inc = reinterpret_cast<const UInt8 *>(input);
	outc = reinterpret_cast<UInt8 *>(tempspace);

	for(i=0;i<count;i++)
	{
		n = arr1[i];
		if(!n)
			n = count-1;
		else
			n--;

		if(!n)
			pi = i;

		outc[i] = inc[n];
	}

	*pindex = pi;

	return outc;
}


#define FLUSH_RLE	\
	do \
	{ \
		if(inRun) \
		{ \
			runLength++; \
			bits = 0;\
			for(int i=1;i<30;i++) \
				if(runLength & (1 << i)) bits = i;\
			while(bits)\
			{\
				bits--;\
				if(runLength & (1 << bits))\
					*output = 1;\
				else\
					*output = 0;\
				output++;\
			}\
			runLength = 0;\
			inRun = 0;\
		} \
	} while(0)

// RC5 uses a modified MTF that doesn't shift on 1
inline UInt8 WriteMTF(UInt8 original, UInt8 *indexes, UInt8 last)
{
	UInt8 result;

	result = indexes[original];

	if((MTFBehavior_k == MTFB_NoSortOn1_k || MTFBehavior_k == MTFB_NoSortOn1LastNonZero_k) && result == 1)
	{
		if(MTFBehavior_k == MTFB_NoSortOn1_k || last != 0)
		{
			for(int i=0;i<256;i++)
			{
				if(!indexes[i])
				{
					indexes[i] = 1;
					break;
				}
			}
			indexes[original] = 0;
		}
	}
	else if(result)
	{
		for(int i=0;i<256;i++)
		{
			if(indexes[i] && indexes[i] < result)
				indexes[i]++;
		}
		indexes[original] = 1;
	}

	return result;
}

// MTF/RLE-0
static const UInt16 *MTFRLE(const void *input, void *tempspace, UInt32 count, UInt32 *numValues)
{
	UInt8 indexes[256];
	const UInt8 *inc;
	UInt16 *output;
	UInt16 *outbase;
	UInt8 last = 255;
	UInt32 bits;

	int inRun = 0;
	UInt32 runLength = 0;

	// Move to good spots for temp space
	inc = reinterpret_cast<const UInt8 *>(input);

	outbase = reinterpret_cast<UInt16 *>(tempspace);
	outbase += (count*2);

	output = outbase;

	// Set up MTF coder
	for(int i=0;i<256;i++)
		indexes[i] = static_cast<UInt8>(i);

	while(count)
	{
		last = WriteMTF(*inc, indexes, last);

		if(last)
		{
			FLUSH_RLE;
			*output = ((unsigned short)last) + 1;
			output++;
		}
		else
		{
			inRun = 1;
			runLength++;
		}

		inc++;
		count--;
	}

	FLUSH_RLE;

	*numValues = output - outbase;

	return outbase;
}

// Returns number of bytes encoded
inline int EncodeSymbol(RangeCoder::Encoder *rc, const SimpleProbabilityTable *ptable, unsigned int sym, UInt8 *output)
{
	UInt32 low = 0;
	for(unsigned int i=0;i<sym;i++)
		low += ptable->probabilities[i];
	rc->Encode(low, ptable->probabilities[sym], ptable->probTotal);
	int nOut = rc->GetOutputCount();
	const UInt8 *rcOut = rc->GetOutput();
	for(int i=0;i<nOut;i++)
		output[i] = rcOut[i];
	return nOut;
}

inline UInt32 WriteSymbol(RangeCoder::Encoder *rc, ProbabilityPool *mset, UInt16 sym, UInt8 *output)
{
	unsigned char *outBase;
	unsigned int firstSym;
	unsigned int secondModel = 0;
	unsigned int secondSym;

	SimpleProbabilityTable ptable;

	outBase = output;

	if(sym == 0) firstSym = ValRun0_k;
	else if(sym == 1) firstSym = ValRun1_k;
	else if(sym == 2) firstSym = Val1_k;
	else
	{
		sym--;
		if(sym & 0x80)			{ firstSym = Val128_255_k;	secondModel = M_128_255_k;	secondSym = sym - 128; }
		else if(sym & 0x40)		{ firstSym = Val64_127_k;	secondModel = M_64_127_k;	secondSym = sym - 64; }
		else if(sym & 0x20)		{ firstSym = Val32_63_k;	secondModel = M_32_63_k;	secondSym = sym - 32; }
		else if(sym & 0x10)		{ firstSym = Val16_31_k;	secondModel = M_16_31_k;	secondSym = sym - 16; }
		else if(sym & 0x08)		{ firstSym = Val8_15_k;		secondModel = M_8_15_k;		secondSym = sym - 8; }
		else if(sym & 0x04)		{ firstSym = Val4_7_k;		secondModel = M_4_7_k;		secondSym = sym - 4; }
		else /*if(sym & 0x02)*/	{ firstSym = Val2_3_k;		secondModel = M_2_3_k;		secondSym = sym - 2; }
	}

	ptable.probabilities = mset->probPool;
	ptable.probTotal = mset->ptotals[0];
	output += EncodeSymbol(rc, &ptable, firstSym, output);
	mset->UpdateProbabilities(0, firstSym);

	if(secondModel)
	{
		ptable.probabilities = mset->probPool + ProbPoolOffsets_k[secondModel];
		ptable.probTotal = mset->ptotals[secondModel];
		output += EncodeSymbol(rc, &ptable, secondSym, output);
		mset->UpdateProbabilities(secondModel, secondSym);
	}

	return output - outBase;
}

static UInt32 EncodeBlock(const UInt16 *input, UInt8 *output, UInt32 count, UInt32 limit)
{
	UInt32 t = 0;
	UInt32 written;
	RangeCoder::Encoder rc;
	ProbabilityPool set;

	rc.Init();
	set.Init();

	// Write each value
	while(count)
	{
		written = static_cast<UInt32>(WriteSymbol(&rc, &set, *input, output));

		output += written;
		t += written;
		input++;
		count--;

		if(t >= limit)
			return 0;
	}

	// Flush what's left
	rc.Flush();
	UInt32 nFlushed = static_cast<UInt32>(rc.GetOutputCount());
	memcpy(output, rc.GetOutput(), nFlushed);
	t += nFlushed;

	if(t >= limit)
		return 0;

	return t;
}



UInt32 TDP::Compression::RC5::CompressBlock(const void *input, void *tempspace, UInt32 count, void *ftab)
{
	UInt32 pindex;
	UInt32 rleValues;
	UInt32 finalValue;

	UInt8 *finaloutput;

	const void *bwtoutput;
	const UInt16 *rleoutput;

	bwtoutput = RC5BlockSort(input, tempspace, count, &pindex, static_cast<UInt32*>(ftab));
	rleoutput = MTFRLE(bwtoutput, tempspace, count, &rleValues);

	finaloutput = static_cast<UInt8*>(tempspace);

	finalValue = EncodeBlock(rleoutput, finaloutput+6, rleValues, count-6);

	if(!finalValue)
	{
		memcpy(finaloutput, input, count);
		return count;
	}

	finaloutput[0] = (unsigned char) (pindex >> 16);
	finaloutput[1] = (unsigned char) (pindex >> 8);
	finaloutput[2] = (unsigned char) pindex;
	finaloutput[3] = (unsigned char) (rleValues >> 16);
	finaloutput[4] = (unsigned char) (rleValues >> 8);
	finaloutput[5] = (unsigned char) rleValues;

	return finalValue+6;
}
